#include <iostream>
#include "opencv2/opencv.hpp"

__global__ void rgba_to_greyscale(unsigned char* r, unsigned char* g, unsigned char*  b,
                                unsigned char* grey, int size)
{
    int index = threadIdx.x + blockIdx.x * blockDim.x;
    // grey = 0.299f * R + 0.587f * G + 0.114f * B
    if(index < size){
        grey[index] = 0.299 * r[index] + 0.587 * g[index] + 0.114 * b[index];
        // printf("%u, %u, %u, %u", grey[index], r[index], g[index], b[index]);
    }
}


int main(int argc, char** argv)
{
    std::cout << "The test run for rgb to gray image." << std::endl;
    
    // read image
    auto image = cv::imread("test.jpg");
    std::cout << "the size of image is: " << image.cols << " " << image.rows << std::endl;
    cv::Mat channel[3];
    cv::split(image, channel);

    cv::imwrite("a.jpg", channel[0]);
    cv::imwrite("b.jpg", channel[1]);
    cv::imwrite("c.jpg", channel[2]);
      
    // copy to gpu
    int t_bytes = image.cols * image.rows * sizeof(unsigned char);
    unsigned char* d_channel_r;
    unsigned char* d_channel_g;
    unsigned char* d_channel_b;
    unsigned char* d_grey;
    cudaMalloc((void**)&d_channel_r, t_bytes);
    cudaMalloc((void**)&d_channel_g, t_bytes);
    cudaMalloc((void**)&d_channel_b, t_bytes);
    cudaMalloc((void**)&d_grey, t_bytes);

    cudaMemcpy((void*)d_channel_r, channel[0].data, t_bytes, cudaMemcpyHostToDevice);
    cudaMemcpy((void*)d_channel_g, channel[1].data, t_bytes, cudaMemcpyHostToDevice);
    cudaMemcpy((void*)d_channel_b, channel[2].data, t_bytes, cudaMemcpyHostToDevice);

    // run kernel function
    dim3 Db(1024);
    dim3 Dg((t_bytes / Db.x) + 1);
    std::cout << "the block of grid:" << Dg.x << std::endl;
    std::cout << "the thread of block:" << Db.x <<std::endl;

    rgba_to_greyscale<<<Dg, Db>>>(d_channel_r, d_channel_g, d_channel_b, d_grey, t_bytes);

    // copy to cpu
    cv::Mat grey_image(image.rows, image.cols, CV_8UC1);
    cudaMemcpy(grey_image.data, (void*)d_grey, t_bytes, cudaMemcpyDeviceToHost);

    // write to disk
    cv::imwrite("grey.jpg", grey_image);

    // free memory
    cudaFree(d_channel_r);
    cudaFree(d_channel_g);
    cudaFree(d_channel_b);
    cudaFree(d_grey);
 
    return 0;
}